#ifndef WTFTPFORUBOOT_H
#define WTFTPFORUBOOT_H
#include <boost/asio.hpp>
#include <boost/filesystem.hpp>
#include <vector>
#include <thread>
#include <mutex>
#include <map>
#include <fstream>
#include <iostream>
#include <iomanip>
#include <functional>

namespace asio = boost::asio;
namespace fs = boost::filesystem;
using asio::ip::udp;

class WTftpForUboot {
public:
    struct ProgressInfo {
        std::string clientIP;
        size_t transferredBytes = 0;
        size_t totalBytes = 0;
    };

    WTftpForUboot(asio::io_context& io, int maxClients = 5);
    ~WTftpForUboot();

    bool start(unsigned short port = 69);
    void print_progress() const;

private:
    struct ClientContext {
        udp::endpoint endpoint;
        std::string filename;
        std::shared_ptr<udp::socket> socket;
    };
    
    std::unordered_map<uint16_t, bool> ack_received_;
    std::unordered_map<uint16_t, boost::asio::steady_timer*> timers_;

    // Core functionality
    void start_receive();
    void handle_request(size_t bytes_received);
    void handle_client(const ClientContext& ctx);
    void handle_receive(const boost::system::error_code& ec, size_t bytes);
    void handle_timeout(const boost::system::error_code& ec, uint16_t expected);
    bool send_block(const ClientContext& ctx, uint16_t block, const char* data, size_t size);
    bool wait_ack(const ClientContext& ctx, uint16_t expected);

    // Helpers
    std::string client_id(const ClientContext& ctx) const;
    void update_progress(const ClientContext& ctx, size_t transferred, size_t total);
    void cleanup_threads();
    void send_error(const ClientContext& ctx, uint16_t code, const std::string& msg);
    bool validate_request(size_t bytes_received) const;

    // Resources
    asio::io_context& _io;
    udp::socket _socket;
    std::array<unsigned char, 512> _recvBuffer;
    udp::endpoint _remoteEndpoint;

    // State management
    std::vector<std::thread> _threads;
    mutable std::mutex _mutex;
    std::map<std::string, ProgressInfo> _progressMap;
    const int _maxClients;
    const fs::path _basePath;
    unsigned short _port = 0;

    // File operations
    std::ifstream open_file(const fs::path& path) const;
};

#endif // WTFTPFORUBOOT_H